import numpy as np
import sys
import json
import meshio


# 导入catia dat网格，类似于.nas
def import_dat(filename):
    with open(filename) as file:
        fileStrings = np.array(file.read().split())
    argwhereGRID = np.argwhere(fileStrings == 'GRID*')
    argwhereStar = np.argwhere(fileStrings == '*')
    if argwhereGRID.size == 0 or argwhereStar.size == 0:
        return
    if argwhereStar[0] - argwhereGRID[0] == 5:
        argV = argwhereGRID + [2, 3, 7]
    elif argwhereStar[0] - argwhereGRID[0] == 4:
        argV = argwhereGRID + [2, 3, 5]
    v = fileStrings[argV]
    v = v.astype(float)
    argwhereCTRIA3 = np.argwhere(fileStrings == 'CTRIA3')
    argwhereCTRIA6 = np.argwhere(fileStrings == 'CTRIA6')
    if argwhereCTRIA3.size > 0:
        argF = argwhereCTRIA3 + [3, 4, 5]
    elif argwhereCTRIA6.size > 0:
        argF = argwhereCTRIA6 + [3, 4, 5]
    f = fileStrings[argF]
    f = f.astype(int) - 1
    return v, f


def import_obj(filename):
    with open(filename) as file:
        s = np.array(file.read().split('\n'))
    v = []
    f = []
    for t in s:
        t = np.array(t.split())
        if t.size > 0:
            if t[0] == 'v':
                v.append(t[1:].astype(float))
            if t[0] == 'f':
                if t[1:].size == 3:
                    f.append(t[1:])
                elif t[1:].size == 4:
                    f.append([t[1], t[2], t[3]])
                    f.append([t[1], t[3], t[4]])
    v = np.array(v)
    f = np.array(f)
    for i in f:
        i[0] = i[0].split('/')[0]
        i[1] = i[1].split('/')[0]
        i[2] = i[2].split('/')[0]
    f = f.astype(int) - 1
    return v, f


def import_nas(filename):
    with open(filename) as file:
        fileStrings = np.array(file.read().split())
    argwhereGRID = np.argwhere(fileStrings == 'GRID*')
    argV = argwhereGRID + [3, 4, 7]
    v = fileStrings[argV]
    v = v.astype(float)
    argwhereCTRIA3 = np.argwhere(fileStrings == 'CTRIA3')
    if argwhereCTRIA3.size > 0:
        argF = argwhereCTRIA3 + [3, 4, 5]
    f = fileStrings[argF]
    f = f.astype(int) - 1
    return v, f


def import_stl(filename):
    mesh = meshio.read(filename)
    v = mesh.points
    f = mesh.cells_dict['triangle']
    return v, f


# 保存为obj文件
def export_obj(filename, v, f):
    with open(filename, 'w') as file:
        for i in v:
            file.write('v ' + str(i[0]) + ' ' + str(i[1]) + ' ' + str(i[2]) + '\n')
        for i in f:
            file.write('f ' + str(i[0] + 1) + ' ' + str(i[1] + 1) + ' ' + str(i[2] + 1) + '\n')


def mesh_infomation(v, f):
    a = {}
    a['x'], a['y'], a['z'] = np.ptp(v, 0)
    a['vertices_number'] = v.shape[0]
    a['facets_number'] = f.shape[0]
    return a


def mesh_ratio(v, ratio):
    v *= ratio
    return v


def mesh_rotate_x(v, angle):
    cos = np.cos(angle * np.pi / 180)
    sin = np.sin(angle * np.pi / 180)
    matrix = np.array([[1, 0, 0],
                       [0, cos, sin],
                       [0, -sin, cos]])
    v=v @ matrix
    return v


def mesh_rotate_y(v, angle):
    cos = np.cos(angle * np.pi / 180)
    sin = np.sin(angle * np.pi / 180)
    matrix = np.array([[cos, 0, -sin],
                       [0, 1, 0],
                       [sin, 0, cos]])
    v=v @ matrix
    return v


def mesh_rotate_z(v, angle):
    cos = np.cos(angle * np.pi / 180)
    sin = np.sin(angle * np.pi / 180)
    matrix = np.array([[cos, sin, 0],
                       [-sin, cos, 0],
                       [0, 0, 1]])
    v=v @ matrix
    return v


if __name__ == '__main__':
    a = import_dat(r'cube.dat')
    print(a)
